package main.java;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

public class TreeRebuilder {

    public List<Node> rebuildTree(List<Node> allNodes) {
        // 1. 筛选出所有实体节点 (type == 5)
        List<Node> entities = allNodes.stream()
                .filter(n -> n.type == 5)
                .collect(Collectors.toList());

        // 2. 构建实体单元 ID 到实体的映射
        Map<String, List<Node>> unitToEntities = new HashMap<>();

        for (Node entity : entities) {
            // 向上查找实体单元和实体类型节点
            Node typeNode = findParentByType(entity, allNodes, 8);
            Node unitNode = findParentByType(entity, allNodes, 1);

            if (unitNode == null || typeNode == null) continue;

            String key = unitNode.id + "#" + typeNode.entityType;

            // 记录下来分组
            unitToEntities.computeIfAbsent(key, k -> new ArrayList<>()).add(entity);
        }

        // 3. 构建新树
        Map<String, Node> unitRoots = new HashMap<>();

        for (Map.Entry<String, List<Node>> entry : unitToEntities.entrySet()) {
            String[] parts = entry.getKey().split("#");
            String unitId = parts[0];
            String entityType = parts[1];

            // 获取实体单元节点
            Node unitNode = cloneNode(findById(unitId, allNodes));
            unitNode.children = unitNode.children != null ? unitNode.children : new ArrayList<>();
            unitRoots.putIfAbsent(unitId, unitNode);

            // 构建实体类型节点
            Node typeNode = new Node();
            typeNode.id = unitId + "_" + entityType;
            typeNode.pid = unitId;
            typeNode.type = 8; // 标记为类型节点
            typeNode.entityType = entityType;
            typeNode.children = new ArrayList<>(entry.getValue()); // 加入所有实体

            unitNode.children.add(typeNode);
        }

        return new ArrayList<>(unitRoots.values());
    }

    // 向上查找第一个匹配type的父节点
    private Node findParentByType(Node node, List<Node> all, int targetType) {
        Map<String, Node> map = all.stream().collect(Collectors.toMap(n -> n.id, n -> n));
        Node current = node;
        while (current != null && current.pid != null) {
            Node parent = map.get(current.pid);
            if (parent == null) break;
            if (parent.type == targetType) return parent;
            current = parent;
        }
        return null;
    }

    private Node findById(String id, List<Node> all) {
        return all.stream().filter(n -> n.id.equals(id)).findFirst().orElse(null);
    }

    private Node cloneNode(Node n) {
        Node copy = new Node();
        copy.id = n.id;
        copy.pid = n.pid;
        copy.longcode = n.longcode;
        copy.type = n.type;
        copy.entityType = n.entityType;
        return copy;
    }
}
